#!/usr/bin/env python3
# -*- coding: utf-8 -*-
"""
AIGC-EDU 启动器
一个用于执行start.bat的Windows桌面应用程序
支持GUI界面、进度显示和系统托盘功能
"""

import tkinter as tk
from tkinter import ttk, scrolledtext, messagebox
import subprocess
import threading
import os
import sys
import re
from datetime import datetime
try:
    import pystray
    from PIL import Image, ImageDraw
except ImportError:
    pystray = None
    Image = None
    ImageDraw = None

class AIGCLauncher:
    def __init__(self):
        self.root = tk.Tk()
        self.root.title("AIGC-EDU 启动器")
        self.root.geometry("800x600")
        self.root.minsize(600, 400)
        
        # 设置窗口图标
        try:
            self.root.iconbitmap(default="icon.ico")
        except:
            pass
        
        # 变量
        self.process = None
        self.is_running = False
        self.tray_icon = None
        
        # 创建界面
        self.create_widgets()
        
        # 绑定窗口关闭事件
        self.root.protocol("WM_DELETE_WINDOW", self.on_closing)
        
        # 检查start.bat是否存在
        self.check_start_bat()
    
    def get_project_root(self):
        """获取项目根目录路径，兼容开发环境和打包后的exe"""
        if getattr(sys, 'frozen', False):
            # 如果是打包后的exe，使用exe所在目录的上级目录
            exe_dir = os.path.dirname(sys.executable)
            # 假设exe在项目的run目录下，需要回到上级目录
            return os.path.dirname(exe_dir)
        else:
            # 开发环境，使用当前文件的上级目录
            return os.path.dirname(os.path.dirname(os.path.abspath(__file__)))
    
    def create_widgets(self):
        """创建GUI组件"""
        # 主框架
        main_frame = ttk.Frame(self.root, padding="10")
        main_frame.grid(row=0, column=0, sticky=(tk.W, tk.E, tk.N, tk.S))
        
        # 配置网格权重
        self.root.columnconfigure(0, weight=1)
        self.root.rowconfigure(0, weight=1)
        main_frame.columnconfigure(1, weight=1)
        main_frame.rowconfigure(2, weight=1)
        
        # 标题
        title_label = ttk.Label(main_frame, text="AIGC实训实践平台V2.0 启动器", 
                               font=('Arial', 16, 'bold'))
        title_label.grid(row=0, column=0, columnspan=3, pady=(0, 20))
        
        # 控制按钮框架
        button_frame = ttk.Frame(main_frame)
        button_frame.grid(row=1, column=0, columnspan=3, pady=(0, 10), sticky=(tk.W, tk.E))
        button_frame.columnconfigure(1, weight=1)
        
        # 运行按钮
        self.run_button = ttk.Button(button_frame, text="启动AIGC实训实践平台V2.0 服务", 
                                    command=self.run_start_bat, style='Accent.TButton')
        self.run_button.grid(row=0, column=0, padx=(0, 10))
        
        # 停止按钮
        self.stop_button = ttk.Button(button_frame, text="停止服务", 
                                     command=self.stop_process, state='disabled')
        self.stop_button.grid(row=0, column=1, padx=(0, 10))
        
        # 清空日志按钮
        clear_button = ttk.Button(button_frame, text="清空日志", 
                                 command=self.clear_log)
        clear_button.grid(row=0, column=2, padx=(0, 10))
        
        # 最小化到托盘按钮
        if pystray:
            tray_button = ttk.Button(button_frame, text="最小化到托盘", 
                                   command=self.minimize_to_tray)
            tray_button.grid(row=0, column=3)
        
        # 进度条
        self.progress = ttk.Progressbar(main_frame, mode='indeterminate')
        self.progress.grid(row=2, column=0, columnspan=3, sticky=(tk.W, tk.E), pady=(0, 10))
        
        # 状态标签
        self.status_label = ttk.Label(main_frame, text="就绪")
        self.status_label.grid(row=3, column=0, columnspan=3, sticky=tk.W)
        
        # 日志输出区域
        log_frame = ttk.LabelFrame(main_frame, text="执行日志", padding="5")
        log_frame.grid(row=4, column=0, columnspan=3, sticky=(tk.W, tk.E, tk.N, tk.S), pady=(10, 0))
        log_frame.columnconfigure(0, weight=1)
        log_frame.rowconfigure(0, weight=1)
        
        self.log_text = scrolledtext.ScrolledText(log_frame, height=15, width=80)
        self.log_text.grid(row=0, column=0, sticky=(tk.W, tk.E, tk.N, tk.S))
        
        # 添加初始日志
        self.log_message("AIGC实训实践平台V2.0 启动器已就绪")
    
    def check_start_bat(self):
        """检查pm2.config.js文件是否存在"""
        config_path = os.path.join(self.get_project_root(), "pm2.config.js")
        if not os.path.exists(config_path):
            self.log_message(f"警告: 未找到 pm2.config.js 文件: {config_path}", "warning")
            self.run_button.config(state='disabled')
        else:
            self.log_message(f"找到 pm2.config.js 文件: {config_path}")
    
    def clean_ansi_codes(self, text):
        """清理ANSI控制字符和颜色代码"""
        # 移除ANSI颜色代码
        ansi_escape = re.compile(r'\x1B(?:[@-Z\\-_]|\[[0-?]*[ -/]*[@-~])')
        text = ansi_escape.sub('', text)
        
        # 移除常见的ANSI代码格式
        text = re.sub(r'\[\d+m', '', text)
        text = re.sub(r'\[\d+;\d+m', '', text)
        text = re.sub(r'\[\d+;\d+;\d+m', '', text)
        
        return text.strip()
    
    def format_pm2_output(self, message):
        """格式化PM2输出，使其更友好"""
        # 清理ANSI代码
        clean_msg = self.clean_ansi_codes(message)
        
        # 如果是PM2相关的消息，进行特殊处理
        if '[PM2]' in clean_msg:
            if 'Applying action restartProcessId' in clean_msg:
                return "正在重启应用..."
            elif '✓' in clean_msg:
                return "应用重启成功"
            elif 'online' in clean_msg and '│' in clean_msg:
                # 这是状态表格的一行，提取关键信息
                parts = [part.strip() for part in clean_msg.split('│') if part.strip()]
                if len(parts) >= 9 and parts[1] != 'name':  # 跳过表头
                    name = parts[1]
                    status = parts[8]
                    cpu = parts[9]
                    mem = parts[10]
                    return f"应用 {name} 状态: {status}, CPU: {cpu}, 内存: {mem}"
            elif '┌' in clean_msg or '├' in clean_msg or '└' in clean_msg or '│ id │' in clean_msg:
                # 跳过表格边框和表头
                return None
        
        # 如果消息为空或只包含特殊字符，跳过
        if not clean_msg or clean_msg.isspace() or all(c in '┌┐└┘├┤┬┴┼─│ ' for c in clean_msg):
            return None
            
        return clean_msg
    
    def log_message(self, message, level="info"):
        """添加日志消息"""
        # 格式化PM2输出
        formatted_msg = self.format_pm2_output(message)
        if formatted_msg is None:
            return  # 跳过不需要显示的消息
            
        timestamp = datetime.now().strftime("%Y-%m-%d %H:%M:%S")
        
        # 根据日志级别设置颜色
        if level == "error":
            color = "red"
        elif level == "warning":
            color = "orange"
        elif level == "success":
            color = "green"
        else:
            color = "black"
        
        # 格式化消息
        formatted_message = f"[{timestamp}] {formatted_msg}\n"
        
        # 插入到文本框
        self.log_text.config(state='normal')
        self.log_text.insert(tk.END, formatted_message, level)
        self.log_text.config(state='disabled')
        
        # 自动滚动到底部
        self.log_text.see(tk.END)
        
        # 配置标签颜色
        self.log_text.tag_config("error", foreground="red")
        self.log_text.tag_config("success", foreground="green")
        self.log_text.tag_config("warning", foreground="orange")
        self.log_text.tag_config("info", foreground="black")
        
        # 更新界面
        self.root.update_idletasks()
    
    def check_pm2_installed(self):
        """检查PM2是否已安装"""
        try:
            # 在 Windows 上，需要使用 shell=True 来正确执行 .cmd 文件
            result = subprocess.run(["pm2", "--version"], capture_output=True, text=True, timeout=10, shell=True)
            return result.returncode == 0
        except (FileNotFoundError, subprocess.TimeoutExpired):
            return False
        except Exception:
            return False

    def run_start_bat(self):
        """运行pm2命令"""
        if self.is_running:
            messagebox.showwarning("警告", "程序正在运行中，请先停止当前进程")
            return

        # 检查PM2是否安装
        if not self.check_pm2_installed():
            messagebox.showerror(
                "PM2 未安装", 
                "PM2 未安装或不在系统 PATH 中。\n\n" +
                "请先安装 PM2:\n" +
                "1. 确保已安装 Node.js\n" +
                "2. 运行命令: npm install -g pm2\n" +
                "3. 重启命令行窗口\n\n" +
                "详细安装说明请参考 README.md 文件。"
            )
            return

        # 查找pm2.config.js文件
        config_path = os.path.join(self.get_project_root(), "pm2.config.js")

        if not os.path.exists(config_path):
            messagebox.showerror("错误", f"未找到 pm2.config.js 文件: {config_path}")
            return

        # 构建pm2命令
        pm2_command = ["pm2", "start", config_path]

        # 在新线程中运行
        thread = threading.Thread(target=self._run_process, args=(pm2_command,))
        thread.daemon = True
        thread.start()
    
    def _run_process(self, command):
        """在后台线程中运行进程"""
        try:
            self.is_running = True
            
            # 更新UI状态
            self.root.after(0, self._update_ui_running)
            
            self.log_message(f"开始执行: {' '.join(command)}")
            
            # 启动进程 - 在 Windows 上使用 shell=True 来正确执行 .cmd 文件
            self.process = subprocess.Popen(
                command,
                stdout=subprocess.PIPE,
                stderr=subprocess.STDOUT,
                text=True,
                bufsize=1,
                universal_newlines=True,
                shell=True,
                encoding='utf-8',
                errors='replace',
                cwd=self.get_project_root()
            )
            
            # 实时读取输出
            for line in iter(self.process.stdout.readline, ''):
                if line:
                    self.root.after(0, self.log_message, line.strip())
            
            # 等待进程完成
            return_code = self.process.wait()
            
            if return_code == 0:
                self.root.after(0, self.log_message, "执行完成 等待自动弹出窗口 http://localhost:3000", "success")
                # PM2启动成功，更新为服务运行状态
                self.root.after(0, self._update_ui_service_running)
            else:
                self.root.after(0, self.log_message, f"执行失败，退出代码: {return_code}", "error")
                # 启动失败，恢复到停止状态
                self.root.after(0, self._update_ui_stopped)
                
        except Exception as e:
            self.root.after(0, self.log_message, f"执行出错: {str(e)}", "error")
            # 出现异常，恢复到停止状态
            self.root.after(0, self._update_ui_stopped)
        
        finally:
            self.is_running = False
            self.process = None
    
    def _update_ui_running(self):
        """更新UI为运行状态"""
        self.run_button.config(state='disabled')
        self.stop_button.config(state='normal')
        self.status_label.config(text="正在运行...")
        self.progress.start()
    
    def _update_ui_stopped(self):
        """更新UI为停止状态"""
        self.run_button.config(state='normal')
        self.stop_button.config(state='disabled')
        self.status_label.config(text="就绪")
        self.progress.stop()
    
    def _update_ui_service_running(self):
        """更新UI为服务运行状态"""
        self.run_button.config(state='normal')  # 保持启动按钮可点击
        self.stop_button.config(state='normal')  # 停止按钮也可点击
        self.status_label.config(text="AIGC服务运行中")
        self.progress.stop()
    
    def stop_process(self):
        """停止PM2进程"""
        if not self.is_running:
            return
            
        try:
            # 执行 pm2 stop all 命令
            self.log_message("正在停止所有PM2进程...", "warning")
            
            result = subprocess.run(
                ["pm2", "stop", "all"], 
                capture_output=True, 
                text=True, 
                shell=True,
                encoding='utf-8',
                errors='replace'
            )
            
            if result.returncode == 0:
                self.log_message("PM2进程已停止", "success")
                # 输出PM2命令的结果
                if result.stdout:
                    for line in result.stdout.strip().split('\n'):
                        if line.strip():
                            self.log_message(line)
            else:
                self.log_message(f"停止PM2进程失败: {result.stderr}", "error")
                
            # 如果有subprocess进程在运行，也要停止它
            if self.process:
                try:
                    self.process.terminate()
                    try:
                        self.process.wait(timeout=3)
                    except subprocess.TimeoutExpired:
                        self.process.kill()
                except Exception:
                    pass
                    
        except Exception as e:
            self.log_message(f"停止进程时出错: {str(e)}", "error")
        
        # 更新UI状态
        self._update_ui_stopped()
    
    def clear_log(self):
        """清空日志"""
        self.log_text.delete(1.0, tk.END)
        self.log_message("日志已清空")
    
    def create_tray_icon(self):
        """创建系统托盘图标"""
        if not pystray or not Image:
            return None
        
        # 创建简单的图标
        image = Image.new('RGB', (64, 64), color='blue')
        draw = ImageDraw.Draw(image)
        draw.rectangle([16, 16, 48, 48], fill='white')
        draw.text((20, 25), "AI", fill='blue')
        
        # 创建托盘菜单
        menu = pystray.Menu(
            pystray.MenuItem("显示窗口", self.show_window),
            pystray.MenuItem("启动 PM2 服务", self.run_start_bat),
            pystray.Menu.SEPARATOR,
            pystray.MenuItem("退出", self.quit_app)
        )
        
        return pystray.Icon("AIGC-EDU", image, menu=menu)
    
    def minimize_to_tray(self):
        """最小化到系统托盘"""
        if not pystray:
            messagebox.showinfo("提示", "系统托盘功能需要安装 pystray 和 Pillow 库")
            return
        
        self.root.withdraw()  # 隐藏窗口
        
        if not self.tray_icon:
            self.tray_icon = self.create_tray_icon()
        
        if self.tray_icon:
            # 在新线程中运行托盘图标
            tray_thread = threading.Thread(target=self.tray_icon.run)
            tray_thread.daemon = True
            tray_thread.start()
    
    def show_window(self, icon=None, item=None):
        """显示窗口"""
        self.root.deiconify()  # 显示窗口
        self.root.lift()  # 置顶
        if self.tray_icon:
            self.tray_icon.stop()
            self.tray_icon = None
    
    def on_closing(self):
        """窗口关闭事件处理"""
        if self.is_running:
            result = messagebox.askyesnocancel(
                "确认退出", 
                "程序正在运行中，是否要停止进程并退出？\n\n" +
                "是：停止进程并退出\n" +
                "否：最小化到托盘\n" +
                "取消：继续运行"
            )
            
            if result is True:  # 是
                self.stop_process()
                self.quit_app()
            elif result is False:  # 否
                self.minimize_to_tray()
            # 取消则什么都不做
        else:
            self.quit_app()
    
    def quit_app(self, icon=None, item=None):
        """退出应用程序"""
        if self.is_running:
            self.stop_process()
        
        if self.tray_icon:
            self.tray_icon.stop()
        
        self.root.quit()
        self.root.destroy()
    
    def run(self):
        """运行应用程序"""
        try:
            self.root.mainloop()
        except KeyboardInterrupt:
            self.quit_app()

def main():
    """主函数"""
    try:
        app = AIGCLauncher()
        app.run()
    except Exception as e:
        print(f"启动应用程序时出错: {e}")
        input("按回车键退出...")

if __name__ == "__main__":
    main()